import cv2
import numpy as np
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QHBoxLayout, QSlider, QLabel, QPushButton, QFrame, QSizePolicy, QScrollArea
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtCore import Qt, pyqtSignal
from core.utils.image_utils import cv2_to_qpixmap  # new import

def center_crop_resize(img, target_width, target_height):
    """
    先按比例缩放，使得缩放后的图像至少覆盖目标尺寸，
    然后从中心裁剪得到目标尺寸的图像.
    """
    orig_height, orig_width = img.shape[:2]
    scale = max(target_width / orig_width, target_height / orig_height)
    resized = cv2.resize(img, (int(orig_width * scale), int(orig_height * scale)))
    
    new_height, new_width = resized.shape[:2]
    x_start = (new_width - target_width) // 2
    y_start = (new_height - target_height) // 2
    cropped = resized[y_start: y_start + target_height, x_start: x_start + target_width]
    return cropped

class AlignmentWidget(QWidget):  # renamed from ManualAlignmentGUI
    offsetChanged = pyqtSignal(int, int)  # new signal
    
    def __init__(self):
        super().__init__()
        # 新增图像列表和索引变量
        self.rgb_image_list = []
        self.ti_image_list = []
        self.image_index = 0
        self.scale_factor = 1.0  # new: store current zoom scale
        self.init_ui()
        
    def init_ui(self):
        self.setWindowTitle("手动对齐工具")
        self.setGeometry(100, 100, 800, 600)
        
        # Main horizontal layout: control panel on left, display area on right.
        main_layout = QHBoxLayout()
        
        # Left: Control Panel wrapped in a fixed-width widget
        self.control_widget = QWidget()
        self.control_widget.setFixedWidth(300)
        control_panel = QVBoxLayout()
        
        # Helper to create slider row with a label for its value.
        def create_slider_row(name, slider, init_value):
            row = QHBoxLayout()
            label = QLabel(name)
            # Update transparency label to use "TI透明度" for the TI slider.
            if name in ["RGB透明度", "TI透明度"]:
                label.setFixedWidth(80)
                slider.setFixedWidth(180)  # extend slider width
            slider.setValue(init_value)
            slider.setToolTip(name)
            value_label = QLabel(str(init_value))
            slider.valueChanged.connect(lambda val, l=value_label: l.setText(str(val)))
            row.addWidget(label)
            row.addWidget(slider)
            row.addWidget(value_label)
            return row, value_label

        # For offset sliders, now store the value labels.
        offset_row, self.offset_x_value_label = create_slider_row("X偏移", QSlider(Qt.Horizontal), 0)
        self.offset_slider_x = offset_row.itemAt(1).widget()  # assuming slider is second widget
        self.offset_slider_x.setRange(-100, 100)
        self.offset_slider_x.valueChanged.connect(self.on_offset_changed)
        
        offset_row2, self.offset_y_value_label = create_slider_row("Y偏移", QSlider(Qt.Horizontal), 0)
        self.offset_slider_y = offset_row2.itemAt(1).widget()
        self.offset_slider_y.setRange(-100, 100)
        self.offset_slider_y.valueChanged.connect(self.on_offset_changed)

        self.rgb_alpha_slider = QSlider(Qt.Horizontal)
        rgb_row, _ = create_slider_row("RGB透明度", self.rgb_alpha_slider, 50)
        self.rgb_alpha_slider.setRange(0, 100)
        self.rgb_alpha_slider.valueChanged.connect(self.update_alignment)
        
        self.thermal_alpha_slider = QSlider(Qt.Horizontal)
        ti_row, _ = create_slider_row("TI透明度", self.thermal_alpha_slider, 50)
        self.thermal_alpha_slider.setRange(0, 100)
        self.thermal_alpha_slider.valueChanged.connect(self.update_alignment)
        
        # Assemble control panel
        control_panel.addLayout(offset_row)
        control_panel.addLayout(offset_row2)
        control_panel.addLayout(rgb_row)  # RGB transparency row.
        control_panel.addLayout(ti_row)   # TI transparency row.
        
        # New buttons with minor style improvements.
        self.toggle_button = QPushButton("隐藏显示区域")
        self.toggle_button.setStyleSheet("QPushButton { padding: 5px; }")
        self.toggle_button.clicked.connect(self.toggle_display)
        
        self.auto_align_button = QPushButton("自动配准")
        self.auto_align_button.setStyleSheet("QPushButton { padding: 5px; }")
        self.auto_align_button.clicked.connect(self.auto_align)
        
        self.restore_button = QPushButton("恢复默认设置")
        self.restore_button.setStyleSheet("QPushButton { padding: 5px; }")
        self.restore_button.clicked.connect(self.restore_defaults)
        
        control_panel.addWidget(self.toggle_button)
        control_panel.addWidget(self.auto_align_button)
        control_panel.addWidget(self.restore_button)
        
        # New: Meta info panel moved into left control panel.
        self.meta_info = QFrame()
        meta_layout = QHBoxLayout()
        self.meta_label1 = QLabel("RGB: N/A")
        self.meta_label2 = QLabel("TI: N/A")
        meta_layout.addWidget(self.meta_label1)
        meta_layout.addWidget(self.meta_label2)
        self.meta_info.setLayout(meta_layout)
        control_panel.addWidget(self.meta_info)
        
        control_panel.addStretch(1)
        self.control_widget.setLayout(control_panel)
        
        # Right: Display Area with annotation support.
        right_panel = QVBoxLayout()
        self.diff_label = QLabel("图像均值绝对差：")
        self.image_label = QLabel()
        self.image_label.setAlignment(Qt.AlignCenter)
        self.scroll_area = QScrollArea()
        self.scroll_area.setWidget(self.image_label)
        self.scroll_area.setWidgetResizable(True)
        right_panel.addWidget(self.diff_label)
        right_panel.addWidget(self.scroll_area)
        
        # Add both panels to main layout
        main_layout.addWidget(self.control_widget)
        main_layout.addLayout(right_panel)
        self.setLayout(main_layout)
        
        # Optional: add a style sheet for overall UI beautification.
        self.setStyleSheet("""
            QLabel { font-size: 14px; }
            QSlider { margin: 5px; }
            QWidget { background-color: #f0f0f0; }
        """)
        
        # New: add navigation and zoom buttons
        nav_layout = QHBoxLayout()
        btn_zoom_in = QPushButton("放大")
        btn_zoom_in.clicked.connect(self.zoom_in)
        btn_zoom_out = QPushButton("缩小")
        btn_zoom_out.clicked.connect(self.zoom_out)
        btn_prev = QPushButton("上一张")
        btn_prev.clicked.connect(self.prev_image)
        btn_next = QPushButton("下一张")
        btn_next.clicked.connect(self.next_image)
        nav_layout.addWidget(btn_zoom_in)
        nav_layout.addWidget(btn_zoom_out)
        nav_layout.addWidget(btn_prev)
        nav_layout.addWidget(btn_next)
        control_panel.addLayout(nav_layout)
        
    def toggle_display(self):
        """Hide or show the display area."""
        # Assuming the second layout in main_layout holds both diff_label and image_label.
        visible = self.image_label.isVisible()
        self.image_label.setVisible(not visible)
        self.diff_label.setVisible(not visible)
        self.toggle_button.setText("显示显示区域" if visible else "隐藏显示区域")
        
    def auto_align(self):
        """Perform automatic alignment (dummy algorithm)."""
        # Example: set predetermined offsets; replace with real logic as needed.
        self.offset_slider_x.setValue(-10)
        self.offset_slider_y.setValue(-10)
        self.update_alignment()
        
    def restore_defaults(self):
        """Restore default settings."""
        self.offset_slider_x.setValue(0)
        self.offset_slider_y.setValue(0)
        self.rgb_alpha_slider.setValue(50)
        self.thermal_alpha_slider.setValue(50)
        self.update_alignment()
        
    def load_images(self, rgb_path, thermal_path):
        """加载图像，并将RGB图像中心裁剪缩放到Thermal尺寸"""
        self.rgb = cv2.imread(rgb_path)
        self.thermal = cv2.imread(thermal_path)
        if self.rgb is None:
            raise ValueError(f"无法加载RGB图像: {rgb_path}")
        if self.thermal is None:
            raise ValueError(f"无法加载Thermal图像: {thermal_path}")
        
        target_width, target_height = self.thermal.shape[1], self.thermal.shape[0]
        self.rgb_resized = center_crop_resize(self.rgb, target_width, target_height)
        self.aligned_rgb = self.rgb_resized.copy()
        self.update_alignment()
        # Remove the forcing resize to avoid geometry errors on small displays:
        # self.resize(300 + target_width, target_height)
        self.meta_label1.setText(f"RGB: {self.rgb.shape[1]}x{self.rgb.shape[0]}")
        self.meta_label2.setText(f"TI: {self.thermal.shape[1]}x{self.thermal.shape[0]}")
        
    def update_alignment(self):
        """根据控制参数进行图像叠加并计算图像差异"""
        offset_x = self.offset_slider_x.value()
        offset_y = self.offset_slider_y.value()
        M = np.float32([[1, 0, offset_x], [0, 1, offset_y]])
        self.aligned_rgb = cv2.warpAffine(self.rgb_resized, M, (self.thermal.shape[1], self.thermal.shape[0]))
        self.show_results()
        
    def on_offset_changed(self, val):
        # Whenever either slider changes, emit the offsetChanged signal.
        x = self.offset_slider_x.value()
        y = self.offset_slider_y.value()
        self.offsetChanged.emit(x, y)
        self.update_alignment()  # update display as before.
        
    def show_results(self):
        """显示叠加图像和差异指标"""
        rgb_weight = self.rgb_alpha_slider.value() / 100.0
        thermal_weight = self.thermal_alpha_slider.value() / 100.0
        overlay = cv2.addWeighted(self.aligned_rgb, rgb_weight, self.thermal, thermal_weight, 0)
        diff = cv2.absdiff(self.aligned_rgb, self.thermal)
        mean_diff = float(np.mean(diff))
        self.diff_label.setText(f"图像均值绝对差：{mean_diff:.2f}")
        # Save original pixmap for zoom purposes.
        self.original_pixmap = cv2_to_qpixmap(overlay)
        self.update_display()
    
    def update_display(self):
        """Render the stored original_pixmap scaled by the current factor."""
        if hasattr(self, 'original_pixmap'):
            w = int(self.original_pixmap.width() * self.scale_factor)
            h = int(self.original_pixmap.height() * self.scale_factor)
            scaled = self.original_pixmap.scaled(w, h, Qt.KeepAspectRatio, Qt.SmoothTransformation)
            self.image_label.setPixmap(scaled)
    
    def zoom_in(self):
        self.scale_factor *= 1.25
        self.update_display()
    
    def zoom_out(self):
        self.scale_factor /= 1.25
        self.update_display()
    
    # 新增方法：设置图像列表并加载第一帧
    def set_images(self, rgb_list, ti_list):
        self.rgb_image_list = rgb_list
        self.ti_image_list = ti_list
        self.image_index = 0
        self.load_current_image()
    
    # 新增方法：加载当前索引的图像
    def load_current_image(self):
        if not self.rgb_image_list or not self.ti_image_list:
            return
        rgb_path = self.rgb_image_list[self.image_index]
        ti_path = self.ti_image_list[self.image_index]
        self.load_images(rgb_path, ti_path)
    
    # 新增方法：加载下一张图像
    def next_image(self):
        if self.rgb_image_list and self.image_index < len(self.rgb_image_list) - 1:
            self.image_index += 1
            self.load_current_image()
    
    # 新增方法：加载上一张图像
    def prev_image(self):
        if self.rgb_image_list and self.image_index > 0:
            self.image_index -= 1
            self.load_current_image()